from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, confusion_matrix
import numpy as np
from .data_utils import dataset_to_numpy

def train_and_evaluate(model, X_train, y_train, X_test, y_test):
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    acc = accuracy_score(y_test, y_pred)
    cm = confusion_matrix(y_test, y_pred)
    return acc, cm, y_pred

def run_all_models(mnist_train, mnist_test, cifar_train, cifar_test, n_samples=5000):
    # MNIST
    X_train_m, y_train_m = dataset_to_numpy(mnist_train)
    X_test_m, y_test_m = dataset_to_numpy(mnist_test)
    # CIFAR-10
    X_train_c, y_train_c = dataset_to_numpy(cifar_train)
    X_test_c, y_test_c = dataset_to_numpy(cifar_test)
    # 只取部分样本以加快实验速度
    X_train_m, y_train_m = X_train_m[:n_samples], y_train_m[:n_samples]
    X_test_m, y_test_m = X_test_m[:1000], y_test_m[:1000]
    X_train_c, y_train_c = X_train_c[:n_samples], y_train_c[:n_samples]
    X_test_c, y_test_c = X_test_c[:1000], y_test_c[:1000]
    # 标准化
    scaler_m = StandardScaler().fit(X_train_m)
    X_train_m = scaler_m.transform(X_train_m)
    X_test_m = scaler_m.transform(X_test_m)
    scaler_c = StandardScaler().fit(X_train_c)
    X_train_c = scaler_c.transform(X_train_c)
    X_test_c = scaler_c.transform(X_test_c)
    # 定义模型
    models = {
        '线性回归': LinearRegression(),
        '逻辑回归': LogisticRegression(max_iter=100),
        'SVM': SVC(),
        '决策树': DecisionTreeClassifier(),
        'KNN': KNeighborsClassifier(),
        'MLP': MLPClassifier(max_iter=20)
    }
    results = {'MNIST': {}, 'CIFAR-10': {}}
    for name, model in models.items():
        acc_m, cm_m, _ = train_and_evaluate(model, X_train_m, y_train_m, X_test_m, y_test_m)
        results['MNIST'][name] = {'acc': acc_m, 'cm': cm_m}
        acc_c, cm_c, _ = train_and_evaluate(model, X_train_c, y_train_c, X_test_c, y_test_c)
        results['CIFAR-10'][name] = {'acc': acc_c, 'cm': cm_c}
        print(f"{name} - MNIST准确率: {acc_m:.3f}, CIFAR-10准确率: {acc_c:.3f}")
    return results
